/*
 * simple_gm_test.cc
 *
 *  Created on: Sep 8, 2011
 *      Author: tqlong
 */

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp>

#include "../simple_gm.h"
#include "../message_passing.h"
#include "../subgradient_parallel.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);
void usage(int exitcode = 0)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

void simple_gm_test()
{
  SimpleGM gm;
  gm.load(vm["gmfile"].as<std::string>().c_str());
  std::cout << "GM loaded ... " << std::endl;
  gm.print_params();

  double Z = gm.partition_function();
  std::cout << "log(Z) = " << Z << " " << exp(Z) << std::endl;

  gm.naive_marginals();
  std::cout << "After computation " << std::endl;
  gm.print_marginals();

  SimpleGM::assignment_type map = gm.naive_map();
  std::cout << "map = (";
  for (int i = 0; i < gm.n(); i++) std::cout << "x" << i << "=" << map[i] << ",";
  std::cout << ")" << std::endl;

  SimpleGM::assignment_type a;
  a[2] = 0;
  SimpleGM rgm = gm.reduce(a);
  std::cout << "Reduced gm " << std::endl;
  rgm.print_params();

  rgm.naive_marginals();
  std::cout << "After computation " << std::endl;
  rgm.print_marginals();

  map = rgm.naive_map();
  std::cout << "map = (";
  for (int i = 0; i < rgm.n(); i++) std::cout << "x" << rgm.new2old_[i] << "=" << map[i] << ",";
  std::cout << ")" << std::endl;
}

void message_passing_test()
{
  std::cout << " ================ MESSAGE PASSING ================ " << std::endl;
  MessagePassingGM gm;
  gm.load(vm["gmfile"].as<std::string>().c_str());
  std::cout << "GM loaded ... " << std::endl;
  gm.print_params();

  double Z = gm.partition_function();
  std::cout << "log(Z) = " << Z << " " << exp(Z) << std::endl;

  gm.message_passing_marginals();
  std::cout << "After computation " << std::endl;
  gm.print_marginals();

  MessagePassingGM::assignment_type map = gm.message_passing_map();
  std::cout << "map = (";
  for (int i = 0; i < gm.n(); i++) std::cout << "x" << i << "=" << map[i] << ",";
  std::cout << ")" << std::endl;

  MessagePassingGM::assignment_type a;
  a[2] = 0;
  MessagePassingGM rgm = gm.reduce(a);
  std::cout << "Reduced gm " << std::endl;
  rgm.print_params();

  rgm.message_passing_marginals();
  std::cout << "After computation " << std::endl;
  rgm.print_marginals();

  map = rgm.message_passing_map();
  std::cout << "map = (";
  for (int i = 0; i < rgm.n(); i++) std::cout << "x" << rgm.new2old_[i] << "=" << map[i] << ",";
  std::cout << ")" << std::endl;
}

void subgradient_test()
{
  std::cout << " ================ SUB-GRADIENT ================ " << std::endl;
  SubGradGM gm;
  gm.load(vm["gmfile"].as<std::string>().c_str());
  std::cout << "GM loaded ... " << std::endl;
  gm.print_params();

  SubGradGM::assignment_type map = gm.subgradient_lp_map();
  std::cout << "map = (";
  for (int i = 0; i < gm.n(); i++) std::cout << "x" << i << "=" << map[i] << ",";
  std::cout << ")" << std::endl;

//  MessagePassingGM::assignment_type a;
//  a[2] = 0;
//  MessagePassingGM rgm = gm.reduce(a);
//  std::cout << "Reduced gm " << std::endl;
//  rgm.print_params();
//
//  rgm.message_passing_marginals();
//  std::cout << "After computation " << std::endl;
//  rgm.print_marginals();
//
//  map = rgm.message_passing_map();
//  std::cout << "map = (";
//  for (int i = 0; i < rgm.n(); i++) std::cout << "x" << rgm.new2old_[i] << "=" << map[i] << ",";
//  std::cout << ")" << std::endl;
}

typedef std::pair<int, int> grid_t;
typedef std::map<grid_t, int> grid2idx_t;
typedef SubGradGM gm_t;
//typedef MessagePassingGM gm_t;
//typedef SimpleGM gm_t;

gm_t noise_gm(int N, grid2idx_t& X, grid2idx_t& Y)
{
  gm_t gm;

  // add nodes
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++)
    {
      arma::vec v(2);
      v[0] = -0.6; v[1] = -0.4;
      grid_t g = grid_t(i,j);

      int idx = gm.addNode(2);
      gm.addPotential(idx, v);
      X[g] = idx;

      idx = gm.addNode(2);
      gm.addPotential(idx, v);
      Y[g] = idx;
    }

  // add edges
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++)
    {
      arma::mat v(2,2);
      v(0,0) = 0; v(0,1) = -5.0;
      v(1,0) = -5.0; v(1,1) = 0;
      grid_t g = grid_t(i,j);
      // edge X(i, j) -- Y(i, j)
      gm.addEdge(X[g], Y[g]);
      gm.addPotential(X[g], Y[g], v);

      // add edges Y(i, j) ---  Y(i-1, j) , Y(i+1, j)  ,  Y(i, j-1) , Y(i, j+1)
      v(0,0) = 0; v(0,1) = -2.3;
      v(1,0) = -2.3; v(1,1) = 0;

      if (i > 0)
      {
        gm.addEdge(Y[g], Y[grid_t(i-1,j)]);
        gm.addPotential(Y[g], Y[grid_t(i-1,j)], v);
      }

      if (j > 0)
      {
        gm.addEdge(Y[g], Y[grid_t(i,j-1)]);
        gm.addPotential(Y[g], Y[grid_t(i,j-1)], v);
      }

      if (i < N-1)
      {
        gm.addEdge(Y[g], Y[grid_t(i+1,j)]);
        gm.addPotential(Y[g], Y[grid_t(i+1,j)], v);
      }

      if (j < N-1)
      {
        gm.addEdge(Y[g], Y[grid_t(i,j+1)]);
        gm.addPotential(Y[g], Y[grid_t(i,j+1)], v);
      }
    }
  return gm;
}

gm_t noise_gm1(int N, int K, grid2idx_t& X, grid2idx_t& Y)
{
  gm_t gm;

  // add nodes
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++)
    {
      arma::vec v(K);
      v.fill(-1.0/K);
      grid_t g = grid_t(i,j);

      int idx = gm.addNode(K);
      gm.addPotential(idx, v);
      X[g] = idx;

      idx = gm.addNode(K);
      gm.addPotential(idx, v);
      Y[g] = idx;
    }

  // add edges
  for (int i = 0; i < N; i++)
    for (int j = 0; j < N; j++)
    {
      arma::mat v(K,K);
      for (int k1=0; k1<K; k1++)
        for (int k2=0; k2<K; k2++)
          v(k1,k2) = -5*(k1-k2)*(k1-k2);
      grid_t g = grid_t(i,j);
      // edge X(i, j) -- Y(i, j)
      gm.addEdge(X[g], Y[g]);
      gm.addPotential(X[g], Y[g], v);

      // add edges Y(i, j) ---  Y(i-1, j) , Y(i+1, j)  ,  Y(i, j-1) , Y(i, j+1)
      for (int k1=0; k1<K; k1++)
        for (int k2=0; k2<K; k2++)
          v(k1,k2) = -2*(k1-k2)*(k1-k2);

      if (i > 0)
      {
        gm.addEdge(Y[g], Y[grid_t(i-1,j)]);
        gm.addPotential(Y[g], Y[grid_t(i-1,j)], v);
      }

      if (j > 0)
      {
        gm.addEdge(Y[g], Y[grid_t(i,j-1)]);
        gm.addPotential(Y[g], Y[grid_t(i,j-1)], v);
      }

      if (i < N-1)
      {
        gm.addEdge(Y[g], Y[grid_t(i+1,j)]);
        gm.addPotential(Y[g], Y[grid_t(i+1,j)], v);
      }

      if (j < N-1)
      {
        gm.addEdge(Y[g], Y[grid_t(i,j+1)]);
        gm.addPotential(Y[g], Y[grid_t(i,j+1)], v);
      }
    }
  return gm;
}

void small_noise_reduction()
{
  int N = 10; // grid size
  grid2idx_t X, Y;
  gm_t gm = noise_gm(N, X, Y);

  // random image
  gm_t::assignment_type a;
  for (int i = 0; i < N; i++)
  {
    for (int j = 0; j < N; j++)
    {
      //int val = (i > j) ? 1 : (i < j ? 0 : (arma::randu() < 0.5 ? 0:1));
      //int val = ( i > 3 && i < 6 && j > 2 && j < 10) ? 1 : 0;
      int val = arma::randu() < 0.7 ? 0:1;
      //int val = (i == 5 || i == 6 || i == 7 || j==4 || j == 5) ? 1 : 0;
      a[X[grid_t(i,j)]] = val;
    }
  }

  gm_t rgm = gm.reduce(a); // reduce gm on evidence
  gm_t::assignment_type map = rgm.subgradient_lp_map();
  //gm_t::assignment_type map = rgm.message_passing_map();
  //gm_t::assignment_type map = rgm.naive_map();

  std::cout << " == NOISE ==" << std::endl;
  for (int i = 0; i < N; i++)
  {
    for (int j = 0; j < N; j++)
      std::cout << a[X[grid_t(i,j)]] << " ";
    std::cout << std::endl;
  }

  std::cout << " == IMAGE ==" << std::endl;
  for (int i = 0; i < N; i++)
  {
    for (int j = 0; j < N; j++)
    {
      int oldIdx = Y[grid_t(i,j)];
      int newIdx = rgm.old2new_[oldIdx];
      int val = map[newIdx];
      std::cout << val << " ";
    }
    std::cout << std::endl;
  }
}

int discrete_sample(arma::vec p)
{
  double s = arma::sum(p);
  double r = arma::randu() * s;
  s = 0;
  for (unsigned int i = 0; i < p.n_elem; i++)
  {
    s += p(i);
    if (s >= r) return i;
  }
  return p.n_elem-1;
}

void small_noise_reduction1()
{
  int N = 10; // grid size
  int K = 4;
  grid2idx_t X, Y;
  gm_t gm = noise_gm1(N, K, X, Y);

  // random image
  gm_t::assignment_type a;
  arma::vec r = arma::randu(K);
  r(0) = arma::sum(r)*2;
  for (int i = 0; i < N; i++)
  {
    for (int j = 0; j < N; j++)
    {
      //int val = (i > j) ? 1 : (i < j ? 0 : (arma::randu() < 0.5 ? 0:1));
      //int val = ( i > 3 && i < 6 && j > 2 && j < 10) ? 1 : 0;
      //int val = discrete_sample(r);
      int val = (i == 5 || i == 6 || i == 7 || j==4 || j == 5) ? 3 : 0;
      a[X[grid_t(i,j)]] = val;
    }
  }

  gm_t rgm = gm.reduce(a); // reduce gm on evidence
  gm_t::assignment_type map = rgm.subgradient_lp_map();
  //gm_t::assignment_type map = rgm.message_passing_map();
  //gm_t::assignment_type map = rgm.naive_map();

  std::cout << " == NOISE ==" << std::endl;
  for (int i = 0; i < N; i++)
  {
    for (int j = 0; j < N; j++)
      std::cout << a[X[grid_t(i,j)]] << " ";
    std::cout << std::endl;
  }

  std::cout << " == IMAGE ==" << std::endl;
  for (int i = 0; i < N; i++)
  {
    for (int j = 0; j < N; j++)
    {
      int oldIdx = Y[grid_t(i,j)];
      int newIdx = rgm.old2new_[oldIdx];
      int val = map[newIdx];
      std::cout << val << " ";
    }
    std::cout << std::endl;
  }
}

int main(int argc, char** argv)
{
  std::cout << "simple graphical model tests ..." << std::endl;

  process_options(argc, argv);

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

  if (vm.count("help"))
    usage();
  //simple_gm_test();
  //message_passing_test();
  //subgradient_test();
  //small_noise_reduction();
  small_noise_reduction1();


  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  std::cerr << "Duration = " << duration << "\n";

  usage();
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("gmfile", po::value<std::string>()->default_value("simple1.txt"), "input file (factor graph)")
    //("fin", po::value<std::string>(&fin)->default_value(""), "input file (factor functions)")
    //("out", po::value<std::string>(&out)->default_value(""), "output file")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

